from django.conf import settings
import requests
import logging
from typing import Dict, List

logger = logging.getLogger('apps')


class OCRServiceError(Exception):
    """OCR服务异常"""

    def __init__(self, message, code=None, http_status=500):
        self.message = message
        self.code = code
        self.http_status = http_status
        super().__init__(message)


class OCRService:
    """OCR服务类 使用 Textin API 的PDF转Markdown功能"""

    def __init__(self):
        self.config = settings.AI_CONFIG.get('ocr_service', {})
        self.app_id = self.config.get('app_id', '1a4137efec74924d48c7f79d130340f9')
        self.secret_code = self.config.get('secret_code', 'a1b14d43c208cba6f18f856784f7ed71')
        self.api_url = 'https://api.textin.com/ai/service/v1/pdf_to_markdown'
        self.max_retries = self.config.get('max_retries', 3)
        self.timeout = self.config.get('timeout', 60)  # PDF处理可能需要更长时间

    def recognize_text(self, pdf_file, options: Dict = None) -> str:
        """将PDF文件转换为Markdown格式

        Args:
            pdf_file: PDF文件对象
            options: 转换选项，可包含以下参数：
                - mode: 转换模式，可选值：simple, preserve_layout, preserve_table (默认preserve_layout)
                - language: 语言，默认自动检测

        Returns:
            str: 转换后的Markdown文本
        """
        try:
            # 验证PDF文件
            if not self.validate_pdf(pdf_file):
                raise OCRServiceError("PDF文件验证失败，请检查文件格式和大小", code="INVALID_PDF")

            # 准备请求头
            headers = {
                'x-ti-app-id': self.app_id,
                'x-ti-secret-code': self.secret_code,
                'Content-Type': 'application/octet-stream'
            }

            # 准备URL参数
            params = {}
            mode = 'preserve_layout'  # 默认保留布局

            if options:
                if 'mode' in options and options['mode'] in ['simple', 'preserve_layout', 'preserve_table']:
                    mode = options['mode']
                if 'language' in options:
                    params['language'] = options.get('language')

            params['mode'] = mode

            # 读取PDF内容
            pdf_content = pdf_file.read()

            # 发送请求，带重试机制
            response = None
            for retry in range(self.max_retries):
                try:
                    response = requests.post(
                        self.api_url,
                        headers=headers,
                        params=params,
                        data=pdf_content,
                        timeout=self.timeout
                    )

                    # 检查响应状态
                    if response.status_code == 200:
                        break

                    # 如果是服务器错误，重试
                    if response.status_code >= 500:
                        logger.warning(
                            f"PDF转Markdown API服务器错误，重试 {retry + 1}/{self.max_retries}: {response.status_code}")
                        continue

                    # 其他错误，直接抛出
                    logger.error(f"PDF转Markdown API请求失败: {response.status_code}, {response.text}")
                    raise OCRServiceError(f"PDF转Markdown API请求失败: {response.status_code}",
                                          code=f"HTTP_{response.status_code}")

                except requests.RequestException as e:
                    logger.warning(f"PDF转Markdown API请求异常，重试 {retry + 1}/{self.max_retries}: {str(e)}")
                    if retry == self.max_retries - 1:
                        raise OCRServiceError(f"PDF转Markdown API请求异常: {str(e)}", code="REQUEST_ERROR")

            if not response:
                raise OCRServiceError("PDF转Markdown API请求失败，无响应", code="NO_RESPONSE")

            # 解析响应
            result = response.json()

            # 检查响应码
            if result.get('code') != 200:
                logger.error(f"PDF转Markdown失败: {result.get('message')}")
                raise OCRServiceError(f"PDF转Markdown失败: {result.get('message')}",
                                      code=f"API_{result.get('code')}")

            # 提取Markdown文本
            markdown_text = result.get('result', {}).get('markdown', '')

            # 记录转换结果
            logger.info(f"PDF转Markdown成功，文本长度: {len(markdown_text)}")

            return markdown_text

        except OCRServiceError:
            # 直接抛出已经格式化的错误
            raise
        except Exception as e:
            logger.error(f"PDF转Markdown失败: {str(e)}", exc_info=True)
            raise OCRServiceError(f"PDF转Markdown失败: {str(e)}", code="UNKNOWN_ERROR")

    def recognize_text_from_url(self, pdf_url: str, options: Dict = None) -> str:
        """从URL转换PDF为Markdown

        Args:
            pdf_url: PDF文件URL
            options: 转换选项，同recognize_text

        Returns:
            str: 转换后的Markdown文本
        """
        try:
            # 验证URL
            if not pdf_url or not (pdf_url.startswith('http://') or pdf_url.startswith('https://')):
                raise OCRServiceError("无效的PDF URL", code="INVALID_URL")

            # 准备请求头
            headers = {
                'x-ti-app-id': self.app_id,
                'x-ti-secret-code': self.secret_code,
                'Content-Type': 'text/plain'
            }

            # 准备URL参数
            params = {}
            mode = 'preserve_layout'  # 默认保留布局

            if options:
                if 'mode' in options and options['mode'] in ['simple', 'preserve_layout', 'preserve_table']:
                    mode = options['mode']
                if 'language' in options:
                    params['language'] = options.get('language')

            params['mode'] = mode

            # 发送请求，带重试机制
            response = None
            for retry in range(self.max_retries):
                try:
                    response = requests.post(
                        self.api_url,
                        headers=headers,
                        params=params,
                        data=pdf_url,
                        timeout=self.timeout
                    )

                    # 检查响应状态
                    if response.status_code == 200:
                        break

                    # 如果是服务器错误，重试
                    if response.status_code >= 500:
                        logger.warning(
                            f"PDF URL转Markdown API服务器错误，重试 {retry + 1}/{self.max_retries}: {response.status_code}")
                        continue

                    # 其他错误，直接抛出
                    logger.error(f"PDF URL转Markdown API请求失败: {response.status_code}, {response.text}")
                    raise OCRServiceError(f"PDF URL转Markdown API请求失败: {response.status_code}",
                                          code=f"HTTP_{response.status_code}")

                except requests.RequestException as e:
                    logger.warning(f"PDF URL转Markdown API请求异常，重试 {retry + 1}/{self.max_retries}: {str(e)}")
                    if retry == self.max_retries - 1:
                        raise OCRServiceError(f"PDF URL转Markdown API请求异常: {str(e)}", code="REQUEST_ERROR")

            if not response:
                raise OCRServiceError("PDF URL转Markdown API请求失败，无响应", code="NO_RESPONSE")

            # 解析响应
            result = response.json()

            # 检查响应码
            if result.get('code') != 200:
                logger.error(f"PDF URL转Markdown失败: {result.get('message')}")
                raise OCRServiceError(f"PDF URL转Markdown失败: {result.get('message')}",
                                      code=f"API_{result.get('code')}")

            # 提取Markdown文本
            markdown_text = result.get('result', {}).get('markdown', '')

            # 记录转换结果
            logger.info(f"PDF URL转Markdown成功，文本长度: {len(markdown_text)}")

            return markdown_text

        except OCRServiceError:
            # 直接抛出已经格式化的错误
            raise
        except Exception as e:
            logger.error(f"PDF URL转Markdown失败: {str(e)}", exc_info=True)
            raise OCRServiceError(f"PDF URL转Markdown失败: {str(e)}", code="UNKNOWN_ERROR")

    def validate_pdf(self, pdf_file) -> bool:
        """验证PDF文件

        Args:
            pdf_file: PDF文件对象

        Returns:
            bool: 是否验证通过
        """
        try:
            # 检查文件大小
            max_size = self.config.get('max_file_size', 50 * 1024 * 1024)  # 默认50MB
            if pdf_file.size > max_size:
                logger.warning(f"PDF文件大小超过限制: {pdf_file.size} > {max_size}")
                return False

            # 检查文件格式是否为PDF
            if not pdf_file.name.lower().endswith('.pdf'):
                logger.warning(f"文件格式不是PDF: {pdf_file.name}")
                return False

            return True

        except Exception as e:
            logger.error(f"PDF文件验证失败: {str(e)}")
            return False

    # 保留原来的validate_image方法，以兼容旧代码
    def validate_image(self, image_file) -> bool:
        """验证图片文件(兼容方法)

        Args:
            image_file: 图片文件对象

        Returns:
            bool: 是否验证通过
        """
        try:
            # 检查文件大小
            max_size = self.config.get('max_file_size', 50 * 1024 * 1024)  # 默认50MB
            if image_file.size > max_size:
                logger.warning(f"文件大小超过限制: {image_file.size} > {max_size}")
                return False

            # 如果是PDF，使用PDF验证
            if image_file.name.lower().endswith('.pdf'):
                return True

            # 检查文件格式
            allowed_formats = self.config.get('allowed_formats', ['jpg', 'jpeg', 'png', 'bmp', 'pdf'])
            file_ext = image_file.name.split('.')[-1].lower()
            if file_ext not in allowed_formats:
                logger.warning(f"文件格式不支持: {file_ext}, 支持的格式: {allowed_formats}")
                return False

            return True

        except Exception as e:
            logger.error(f"文件验证失败: {str(e)}")
            return False

    def get_text_structure(self, markdown_text: str) -> Dict:
        """从Markdown文本中提取结构化信息

        Args:
            markdown_text: PDF转Markdown结果文本

        Returns:
            Dict: 结构化信息
        """
        try:
            # 提取标题
            title = self._extract_title_from_markdown(markdown_text)

            # 提取段落
            paragraphs = self._extract_paragraphs_from_markdown(markdown_text)

            # 提取表格
            tables = self._extract_tables_from_markdown(markdown_text)

            # 提取列表
            lists = self._extract_lists_from_markdown(markdown_text)

            return {
                'title': title,
                'paragraphs': paragraphs,
                'tables': tables,
                'lists': lists
            }

        except Exception as e:
            logger.error(f"从Markdown提取结构失败: {str(e)}")
            return {
                'title': '',
                'paragraphs': [],
                'tables': [],
                'lists': []
            }

    def _extract_title_from_markdown(self, markdown_text: str) -> str:
        """从Markdown文本中提取标题

        Args:
            markdown_text: Markdown文本

        Returns:
            str: 标题
        """
        try:
            # 查找第一个一级标题
            lines = markdown_text.split('\n')
            for line in lines:
                line = line.strip()
                if line.startswith('# '):
                    return line[2:].strip()

            # 如果没有一级标题，寻找第一个二级标题
            for line in lines:
                line = line.strip()
                if line.startswith('## '):
                    return line[3:].strip()

            # 如果没有找到标题，返回文档前20个字符
            first_content_line = next((line for line in lines if line.strip()), '')
            return first_content_line[:20] + ('...' if len(first_content_line) > 20 else '')
        except Exception as e:
            logger.error(f"提取Markdown标题失败: {str(e)}")
            return ''

    def _extract_paragraphs_from_markdown(self, markdown_text: str) -> List[str]:
        """从Markdown文本中提取段落

        Args:
            markdown_text: Markdown文本

        Returns:
            List[str]: 段落列表
        """
        try:
            paragraphs = []
            lines = markdown_text.split('\n')
            current_paragraph = []

            for line in lines:
                line = line.strip()
                # 跳过标题、列表项、表格行和分隔线
                if (line.startswith('#') or line.startswith('- ') or line.startswith('* ') or
                        line.startswith('|') or line.startswith('```') or line == '---' or line == '***'):
                    # 如果有积累的段落内容，添加到结果
                    if current_paragraph:
                        paragraphs.append(' '.join(current_paragraph))
                        current_paragraph = []
                    continue

                # 空行表示段落结束
                if not line and current_paragraph:
                    paragraphs.append(' '.join(current_paragraph))
                    current_paragraph = []
                # 非空行且不是特殊元素，添加到当前段落
                elif line:
                    current_paragraph.append(line)

            # 处理最后一个段落
            if current_paragraph:
                paragraphs.append(' '.join(current_paragraph))

            return paragraphs
        except Exception as e:
            logger.error(f"提取Markdown段落失败: {str(e)}")
            return []

    def _extract_tables_from_markdown(self, markdown_text: str) -> List[Dict]:
        """从Markdown文本中提取表格

        Args:
            markdown_text: Markdown文本

        Returns:
            List[Dict]: 表格列表，每个表格包含headers和rows
        """
        try:
            tables = []
            lines = markdown_text.split('\n')

            # 遍历找表格
            i = 0
            while i < len(lines):
                if lines[i].strip().startswith('|') and i + 1 < len(lines) and '---' in lines[i + 1]:
                    # 找到表格头部
                    headers = [cell.strip() for cell in lines[i].strip('|').split('|')]
                    i += 2  # 跳过分隔行

                    # 收集表格行
                    rows = []
                    while i < len(lines) and lines[i].strip().startswith('|'):
                        cells = [cell.strip() for cell in lines[i].strip('|').split('|')]
                        rows.append(cells)
                        i += 1

                    # 添加到表格列表
                    if headers and rows:
                        tables.append({
                            'headers': headers,
                            'rows': rows
                        })
                else:
                    i += 1

            return tables
        except Exception as e:
            logger.error(f"提取Markdown表格失败: {str(e)}")
            return []

    def _extract_lists_from_markdown(self, markdown_text: str) -> List[Dict]:
        """从Markdown文本中提取列表

        Args:
            markdown_text: Markdown文本

        Returns:
            List[Dict]: 列表内容，包含类型和项目
        """
        try:
            lists = []
            lines = markdown_text.split('\n')
            current_list = None

            for i, line in enumerate(lines):
                line = line.strip()

                # 检测列表项
                if line.startswith('- ') or line.startswith('* '):
                    list_item = line[2:].strip()

                    # 如果没有当前列表或上一行不是列表项，创建新列表
                    if not current_list or not (
                            lines[i - 1].strip().startswith('- ') or lines[i - 1].strip().startswith('* ')):
                        if current_list:
                            lists.append(current_list)
                        current_list = {'type': 'unordered', 'items': [list_item]}
                    else:
                        current_list['items'].append(list_item)
                # 检测有序列表
                elif line.match(r'^\d+\.\s'):
                    list_item = line.split('.', 1)[1].strip()

                    # 如果没有当前列表或上一行不是有序列表项，创建新列表
                    if not current_list or not lines[i - 1].strip().match(r'^\d+\.\s'):
                        if current_list:
                            lists.append(current_list)
                        current_list = {'type': 'ordered', 'items': [list_item]}
                    else:
                        current_list['items'].append(list_item)
                # 如果行不是列表项且有当前列表，结束当前列表
                elif current_list:
                    lists.append(current_list)
                    current_list = None

            # 处理最后一个列表
            if current_list:
                lists.append(current_list)

            return lists
        except Exception as e:
            logger.error(f"提取Markdown列表失败: {str(e)}")
            return []

    # 以下方法保留兼容旧代码，但内部逻辑已更新
    def _extract_title(self, result: Dict) -> str:
        """提取标题(兼容方法)"""
        try:
            # 如果是旧格式API结果
            if 'result' in result and 'pages' in result['result']:
                # 使用旧的提取逻辑
                if result['result']['pages']:
                    first_page = result['result']['pages'][0]
                    if 'lines' in first_page and first_page['lines']:
                        return first_page['lines'][0].get('text', '')
            # 如果是新API格式或Markdown文本
            elif 'result' in result and 'markdown' in result['result']:
                return self._extract_title_from_markdown(result['result']['markdown'])

            return ''
        except Exception:
            return ''

    def _extract_paragraphs(self, result: Dict) -> List[str]:
        """提取段落(兼容方法)"""
        try:
            # 如果是旧格式API结果
            if 'result' in result and 'pages' in result['result']:
                # 使用旧的提取逻辑
                paragraphs = []
                for page in result['result']['pages']:
                    if 'lines' in page:
                        current_paragraph = []
                        for line in page['lines']:
                            text = line.get('text', '').strip()
                            if text:
                                current_paragraph.append(text)
                            elif current_paragraph:
                                paragraphs.append(' '.join(current_paragraph))
                                current_paragraph = []

                        if current_paragraph:
                            paragraphs.append(' '.join(current_paragraph))
                return paragraphs
            # 如果是新API格式或Markdown文本
            elif 'result' in result and 'markdown' in result['result']:
                return self._extract_paragraphs_from_markdown(result['result']['markdown'])

            return []
        except Exception:
            return []

    def _extract_tables(self, result: Dict) -> List[Dict]:
        """提取表格(兼容方法)"""
        try:
            # 如果是新API格式或Markdown文本
            if 'result' in result and 'markdown' in result['result']:
                return self._extract_tables_from_markdown(result['result']['markdown'])
            # 旧格式不支持表格，返回空列表
            return []
        except Exception:
            return []

    def _extract_key_info(self, result: Dict) -> Dict:
        """提取关键信息(兼容方法)"""
        # 不再使用，但为保持兼容性而保留
        return {}